#include <stdio.h>
#include <math.h>

#define ABS(a) ((a) > 0 ? (a) : (-(a)))

float din[1][3][32][32];
float dout[1][3][32][32];

float weight[3] = {1.0f,1.0f,1.0f};
float bias[3] = {0.0f,0.0f,0.0f};

void data_pre(){
    FILE *fp = fopen("din.txt", "r");
    for (int i = 0; i < 3*32*32; i++){
        fscanf(fp, "%f", (float *)din + i);
    }
    fclose(fp);
}

void data_export(){
    FILE *fp = fopen("dout_eval.txt", "w");
    for (int i = 0; i < 3*32*32; i++){
        fprintf(fp, "%f\n", *((float *)dout + i));
    }
    fclose(fp);
}

void batch_norm(
    float din[1][3][32][32],
    float weight[3],
    float bias[3],
    float dout[1][3][32][32]
){
    float avg[3] = {0.f};
    float var[3] = {0.f};

    for (int i = 0; i < 3; i++){
        avg[i] = 0.0;
        for (int j = 0; j < 32 * 32; j++){
            avg[i] += *((float *)(din[0][i]) + j);       
        }
        avg[i] /= 32 * 32;
        printf("avg[%d], %f\n", i, avg[i]);
    }

    for (int i = 0; i < 3; i++){
        var[i] = 0.0;
        for (int j = 0; j < 32 * 32; j++){
            float d = *((float *)din[0][i] + j);
            var[i] += (d - avg[i]) * (d - avg[i]);    
        }
        var[i] /= 32 * 32;
        printf("var[%d], %f\n", i, var[i]);
    }

    for (int index_ch = 0; index_ch < 3; index_ch++){
        for (int i = 0; i < 32; i++){
            for (int j = 0; j < 32; j++){
                dout[0][index_ch][i][j] = (din[0][index_ch][i][j] - avg[index_ch]) / sqrt(var[index_ch] + 0.00001);
            }
        }
    }
}

int main(int argc, char const *argv[])
{
    data_pre();
    batch_norm(din, weight, bias, dout);
    data_export();
    return 0;
}
